"""
配置管理器
负责加载和管理系统配置
"""

import yaml
import os
from pathlib import Path
from typing import Dict, Any, Optional
import logging

logger = logging.getLogger(__name__)

class ConfigManager:
    """配置管理器"""
    
    def __init__(self, config_path: Optional[str] = None):
        """初始化配置管理器"""
        self.config_path = config_path or "configs/model_config.yaml"
        self.config = {}
        self._load_config()
    
    def _load_config(self):
        """加载配置文件"""
        try:
            config_file = Path(self.config_path)
            
            if not config_file.exists():
                logger.warning(f"配置文件不存在: {self.config_path}")
                self._create_default_config()
                return
            
            with open(config_file, 'r', encoding='utf-8') as f:
                self.config = yaml.safe_load(f)
            
            logger.info(f"配置加载成功: {self.config_path}")
            
        except Exception as e:
            logger.error(f"加载配置失败: {e}")
            self._create_default_config()
    
    def _create_default_config(self):
        """创建默认配置"""
        logger.info("创建默认配置")
        
        self.config = {
            'text2video': {
                'model_path': 'runwayml/stable-diffusion-v1-5',
                'resolution': [256, 256],
                'fps': 24,
                'num_frames': 16,
                'max_frames': 64,
                'num_inference_steps': 50,
                'guidance_scale': 7.5,
                'seed': 42
            },
            'style_transfer': {
                'model_path': 'models/style_transfer',
                'resolution': [256, 256],
                'styles': {
                    'cartoon': {
                        'name': '卡通风格',
                        'description': '明亮的卡通动画风格'
                    },
                    'anime': {
                        'name': '动漫风格',
                        'description': '日本动漫风格'
                    },
                    'watercolor': {
                        'name': '水彩画风格',
                        'description': '柔和的水彩画效果'
                    }
                }
            },
            'video_editing': {
                'model_path': 'models/video_editing',
                'operations': {
                    'enhance': {
                        'name': '质量增强',
                        'description': '提升视频整体质量'
                    },
                    'stabilize': {
                        'name': '视频稳定',
                        'description': '减少视频抖动'
                    },
                    'denoise': {
                        'name': '降噪',
                        'description': '去除视频噪声'
                    }
                }
            },
            'training': {
                'batch_size': 4,
                'learning_rate': 1e-4,
                'num_epochs': 100,
                'save_interval': 10,
                'eval_interval': 5
            },
            'inference': {
                'batch_size': 1,
                'memory_efficient': True,
                'cache_dir': 'cache',
                'output_format': 'mp4'
            },
            'system': {
                'device': 'auto',
                'log_level': 'INFO',
                'cache_enabled': True
            }
        }
        
        # 保存默认配置
        self.save_config()
    
    def get_config(self) -> Dict[str, Any]:
        """获取配置"""
        return self.config
    
    def get_section(self, section: str) -> Dict[str, Any]:
        """获取配置节"""
        return self.config.get(section, {})
    
    def get_value(self, section: str, key: str, default: Any = None) -> Any:
        """获取配置值"""
        section_config = self.get_section(section)
        return section_config.get(key, default)
    
    def set_value(self, section: str, key: str, value: Any):
        """设置配置值"""
        if section not in self.config:
            self.config[section] = {}
        
        self.config[section][key] = value
        logger.info(f"设置配置: {section}.{key} = {value}")
    
    def update_section(self, section: str, updates: Dict[str, Any]):
        """更新配置节"""
        if section not in self.config:
            self.config[section] = {}
        
        self.config[section].update(updates)
        logger.info(f"更新配置节: {section}")
    
    def save_config(self, config_path: Optional[str] = None):
        """保存配置"""
        try:
            save_path = config_path or self.config_path
            config_file = Path(save_path)
            
            # 确保目录存在
            config_file.parent.mkdir(parents=True, exist_ok=True)
            
            with open(config_file, 'w', encoding='utf-8') as f:
                yaml.dump(self.config, f, default_flow_style=False, 
                         allow_unicode=True, indent=2)
            
            logger.info(f"配置保存成功: {save_path}")
            
        except Exception as e:
            logger.error(f"保存配置失败: {e}")
            raise
    
    def reload_config(self):
        """重新加载配置"""
        logger.info("重新加载配置")
        self._load_config()
    
    def validate_config(self) -> bool:
        """验证配置"""
        try:
            required_sections = ['text2video', 'style_transfer', 'video_editing']
            
            for section in required_sections:
                if section not in self.config:
                    logger.error(f"缺少必需的配置节: {section}")
                    return False
            
            # 验证文本到视频配置
            text2video_config = self.config['text2video']
            required_keys = ['model_path', 'resolution', 'fps', 'num_frames']
            
            for key in required_keys:
                if key not in text2video_config:
                    logger.error(f"缺少必需的配置项: text2video.{key}")
                    return False
            
            logger.info("配置验证通过")
            return True
            
        except Exception as e:
            logger.error(f"配置验证失败: {e}")
            return False
    
    def get_device_config(self) -> str:
        """获取设备配置"""
        device = self.get_value('system', 'device', 'auto')
        
        if device == 'auto':
            import torch
            if torch.cuda.is_available():
                return 'cuda'
            elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                return 'mps'
            else:
                return 'cpu'
        
        return device
    
    def get_model_path(self, model_type: str) -> str:
        """获取模型路径"""
        return self.get_value(model_type, 'model_path', '')
    
    def get_resolution(self, model_type: str) -> list:
        """获取分辨率配置"""
        return self.get_value(model_type, 'resolution', [256, 256])
    
    def get_training_config(self) -> Dict[str, Any]:
        """获取训练配置"""
        return self.get_section('training')
    
    def get_inference_config(self) -> Dict[str, Any]:
        """获取推理配置"""
        return self.get_section('inference')
    
    def get_system_config(self) -> Dict[str, Any]:
        """获取系统配置"""
        return self.get_section('system')
    
    def create_env_config(self, env_path: str = ".env"):
        """创建环境配置文件"""
        try:
            env_content = f"""# AI动画生成系统环境配置

# 模型路径
MODEL_PATH={self.get_model_path('text2video')}
STYLE_MODEL_PATH={self.get_model_path('style_transfer')}
EDIT_MODEL_PATH={self.get_model_path('video_editing')}

# 设备配置
DEVICE={self.get_device_config()}

# 输出配置
OUTPUT_DIR=outputs
CACHE_DIR=cache

# 日志配置
LOG_LEVEL={self.get_value('system', 'log_level', 'INFO')}
LOG_FILE=outputs/logs/system.log

# 训练配置
BATCH_SIZE={self.get_value('training', 'batch_size', 4)}
LEARNING_RATE={self.get_value('training', 'learning_rate', 1e-4)}
NUM_EPOCHS={self.get_value('training', 'num_epochs', 100)}

# 推理配置
INFERENCE_BATCH_SIZE={self.get_value('inference', 'batch_size', 1)}
OUTPUT_FORMAT={self.get_value('inference', 'output_format', 'mp4')}

# 缓存配置
CACHE_ENABLED={self.get_value('system', 'cache_enabled', True)}
CACHE_SIZE={self.get_value('system', 'cache_size', '2GB')}
"""
            
            with open(env_path, 'w', encoding='utf-8') as f:
                f.write(env_content)
            
            logger.info(f"环境配置文件创建成功: {env_path}")
            
        except Exception as e:
            logger.error(f"创建环境配置文件失败: {e}")
            raise
    
    def merge_config(self, other_config: Dict[str, Any]):
        """合并配置"""
        def deep_merge(base: Dict[str, Any], update: Dict[str, Any]) -> Dict[str, Any]:
            """深度合并字典"""
            result = base.copy()
            
            for key, value in update.items():
                if key in result and isinstance(result[key], dict) and isinstance(value, dict):
                    result[key] = deep_merge(result[key], value)
                else:
                    result[key] = value
            
            return result
        
        self.config = deep_merge(self.config, other_config)
        logger.info("配置合并完成")
    
    def export_config(self, export_path: str):
        """导出配置"""
        try:
            with open(export_path, 'w', encoding='utf-8') as f:
                yaml.dump(self.config, f, default_flow_style=False, 
                         allow_unicode=True, indent=2)
            
            logger.info(f"配置导出成功: {export_path}")
            
        except Exception as e:
            logger.error(f"导出配置失败: {e}")
            raise
    
    def import_config(self, import_path: str):
        """导入配置"""
        try:
            with open(import_path, 'r', encoding='utf-8') as f:
                imported_config = yaml.safe_load(f)
            
            self.merge_config(imported_config)
            logger.info(f"配置导入成功: {import_path}")
            
        except Exception as e:
            logger.error(f"导入配置失败: {e}")
            raise

# 全局配置管理器实例
_global_config_manager = None

def get_config_manager() -> ConfigManager:
    """获取全局配置管理器"""
    global _global_config_manager
    
    if _global_config_manager is None:
        _global_config_manager = ConfigManager()
    
    return _global_config_manager

def get_config() -> Dict[str, Any]:
    """获取全局配置"""
    return get_config_manager().get_config()

def get_section(section: str) -> Dict[str, Any]:
    """获取配置节"""
    return get_config_manager().get_section(section)

def get_value(section: str, key: str, default: Any = None) -> Any:
    """获取配置值"""
    return get_config_manager().get_value(section, key, default)

def set_value(section: str, key: str, value: Any):
    """设置配置值"""
    get_config_manager().set_value(section, key, value) 